package com.starhub.application.agent.service.impl;

import com.starhub.application.agent.constants.ModelConstants;
import com.starhub.application.agent.dto.AIChatInfo;
import com.starhub.application.agent.dto.ChatDTO;
import com.starhub.application.agent.enums.ChatMessageTypeEnum;
import com.starhub.application.agent.rag.RagHelper;
import com.starhub.application.agent.result.AIAgentResult;
import com.starhub.application.agent.result.AiChatResult;
import com.starhub.application.agent.service.IAgentAuthService;
import com.starhub.application.agent.service.IChatService;
import com.starhub.application.agent.service.IModelService;
import com.starhub.application.function.biz.AiFunctionProvider;
import com.starhub.application.interfaces.agent.AgentConfigService;
import com.starhub.application.message.dto.Message;
import com.starhub.application.message.service.IMessageStoreService;
import com.starhub.application.message.util.Utils;
import com.starhub.application.model.dto.ModelCallDTO;
import com.starhub.common.bean.agent.AgentConfig;
import com.starhub.common.bean.agent.AiFunction;
import com.starhub.common.bean.agent.PromptVectorConfig;
import com.starhub.common.security.util.SecurityUtils;
import com.starhub.monitor.dto.AIChatLogDto;
import com.starhub.monitor.service.IAIChatLogService;
import com.starhub.utils.Util;

import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.agent.tool.ToolSpecification;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.service.Result;
import dev.langchain4j.service.tool.ToolExecutor;
import lombok.extern.slf4j.Slf4j;

import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.PrintWriter;
import java.time.LocalDateTime;
import java.util.*;
import java.util.Map.Entry;
import java.util.stream.Collectors;

@Slf4j
@Service
public class ChatServiceImpl implements IChatService {

    @Autowired
    private IModelService modelService;

    @Autowired
    private AgentConfigService agentConfigService;

    @Autowired
    private IMessageStoreService messageStoreService;

    @Autowired
    private IAgentAuthService agentAuthService;

    @Autowired
    private AiFunctionProvider aiFunctionProvider;

    @Autowired
    private IAIChatLogService aiChatLogService;

    @Autowired
    private RagHelper ragHelper;

    /**
     * 初始化智能体信息
     */
    @Override
    public AIAgentResult initAgentById(Long agentId) {
        AgentConfig agentConfig = agentConfigService.getById(agentId);
        if (agentConfig == null) {
            return AIAgentResult.fail("未找到指定的智能体");
        }
        return initAgent(agentConfig);
    }

    /**
     * 初始化智能体信息
     */
    @Override
    public AIAgentResult initAgentByMark(String identifier) {
        AgentConfig agentConfig = agentConfigService.getByIdentifier(identifier);
        if (agentConfig == null) {
            return AIAgentResult.fail("未找到指定的智能体");
        }
        return initAgent(agentConfig);
    }

    /**
     * 初始化智能体信息
     * @param agentConfig
     * @return
     */
    public AIAgentResult initAgent(AgentConfig agentConfig) {
        // 1. 智能体权限校验   
        /*LoginUser user = SecurityUtils.getCurrentUser();
        if (!agentAuthService.checkAgentPermission(agentConfig.getId(), user.getUserId())) {
            return AIAgentResult.fail("无权限访问该智能体");
        }*/
        // 2. 生成sessionId
        String sessionId = UUID.randomUUID().toString();

        // 3. 缓存sessionId
        //messageStoreService.saveSessionId(sessionId, agentId);

        // 4. 返回结果
        return AIAgentResult.success(agentConfig, sessionId, "初始化成功");
    }

    /**
     * 执行对话任务
     * 
     * @param chatInfo 对话信息
     * @return 对话结果
     */
    @Override
    public AiChatResult chat(AIChatInfo chatInfo) {
        String userId = SecurityUtils.getCurrentUserId();
        
        // 1. 获取智能体信息
        AgentConfig agentConfig = agentConfigService.getById(chatInfo.getAgentId());
        if (agentConfig == null) {
            return AiChatResult.fail("未找到指定的智能体");
        }

        // 2. 智能体鉴权
        if (!agentAuthService.checkAgentPermission(chatInfo.getAgentId(), userId)) {
            return AiChatResult.fail("无权限访问该智能体");
        }

        // 3. 确定使用的模型标识
        String modelMark =  ModelConstants.MODEL_USAGE_DYNAMIC.equals(agentConfig.getModelUsage()) ? chatInfo.getSelectedModel(): agentConfig.getFixedModel();
        if (modelMark == null) {
            return AiChatResult.fail("未指定模型标识");
        }

        // 4. 获取历史对话
        List<Message> historyMessages = new ArrayList<>();
        if (agentConfig.getReferencedFunctions() != null && !agentConfig.getReferencedFunctions().isEmpty()) {
            historyMessages = messageStoreService.getHistoryMessages(chatInfo.getSessionId(), 100);
        }

        //保存用户消息
        Message userMessage_user = Message.builder()
                .sessionId(chatInfo.getSessionId())
                .userId(userId)
                .content(chatInfo.getUserMessage())
                .role(ChatMessageTypeEnum.USER.getCode())
                .agentId(chatInfo.getAgentId())
                .createdAt(new Date()).build();
        Long chatSessionId = messageStoreService.saveMessage(userMessage_user);

        // 5. 背景提示词
        historyMessages.add(0, Message.builder().content(agentConfig.getBackgroundPrompt())
                .role(ChatMessageTypeEnum.SYSTEM.getCode()).build());

        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(100);

        historyMessages.stream().map(ChatServiceImpl::transMessage).forEach(chatMemory::add);

        // 6. 构建提示词
        String finalPrompt = buildFinalPrompt(chatInfo, agentConfig);

        // 7. 向量知识库检索
        List<PromptVectorConfig> vectorConfigList = agentConfigService.getVectorConfigById(chatInfo.getAgentId());
        if (vectorConfigList != null && !vectorConfigList.isEmpty()) {

            Map<String, String> ragResult = ragHelper.search(chatInfo.getUserMessage(), vectorConfigList);
            if (ragResult != null && !ragResult.isEmpty()) {

                for (Entry<String, String> entry : ragResult.entrySet()) {  
                    String key = entry.getKey();
                    String value = entry.getValue();
                    finalPrompt = finalPrompt.replace("$" + key + "$", value);
                }
            }
        }

        // 8. 处理函数调用
        String[] functionNames = {};
        if (chatInfo.getFunctionName() != null) {
            functionNames = new String[]{chatInfo.getFunctionName()};
        } else if (agentConfig.getReferencedFunctions() != null && !agentConfig.getReferencedFunctions().isEmpty()) {
            // 处理所有可用函数
            functionNames = agentConfig.getReferencedFunctions().split(",");
        }
        List<String> availableFunctions = new ArrayList<>();
        Map<ToolSpecification, ToolExecutor> toolsMap = new HashMap<>();
        for (String functionName : functionNames) {
            AiFunction function = aiFunctionProvider.getFunctionByName(functionName);
            if (function != null) {
                availableFunctions.add(functionName);
                toolsMap.put(Utils.convertToToolSpecification(function), aiFunctionProvider.toolExecutor);
            }
        }

        // 9. 调用模型服务
        Result<String> callResult = modelService.call(new ModelCallDTO(modelMark, agentConfig.getId(),
                finalPrompt, chatMemory, toolsMap, agentConfig.getMaxtoken()));

        //todo: callResult本身的运行时异常       

        String modelResponse = callResult.content();
        
        // 10. 保存模型响应
        Message assistantMessage = Message.builder()
                .sessionId(chatInfo.getSessionId())
                .userId(userId)
                .content(modelResponse)
                .role(ChatMessageTypeEnum.ASSISTANT.getCode())
                .createdAt(new Date()).build();

        messageStoreService.saveMessage(assistantMessage);

        aiChatLogService.logSuccess(AIChatLogDto.builder()
            .agentId(chatInfo.getAgentId())
            .userId(Util.getLongValue(chatInfo.getUserId()))
            .modelName(modelMark)
            .sessionId(chatInfo.getSessionId())
            .messageId(chatSessionId)
            .requestContent(chatInfo.getUserMessage())
            .responseContent(modelResponse)
            .requestTime(chatInfo.getRequestTime())
            .responseTime(LocalDateTime.now())
            .build()
        );

        // 11. 构建返回结果
        AiChatResult result = AiChatResult.success(modelResponse);
        result.setMessageMode(agentConfig.getIsStreaming().equals("1")  ? "stream" : "normal");
        result.setMatchedFunctions(availableFunctions);
        //result.setKnowledgeContent(knowledgeContent);
        result.setChatSessionId(Util.null2String(chatSessionId));

        return result;
    }

    @Override
    @Deprecated
    public SseEmitter chatStream(AIChatInfo chatInfo) {

        // 1. 获取智能体信息
        AgentConfig agentConfig = agentConfigService.getById(chatInfo.getAgentId());
        if (agentConfig == null) {
            return null;
        }

        // 2. 智能体鉴权
        if (!agentAuthService.checkAgentPermission(chatInfo.getAgentId(), chatInfo.getUserId())) {
            return null;
        }

        // 3. 确定使用的模型标识
        String modelMark = ModelConstants.MODEL_USAGE_DYNAMIC.equals(agentConfig.getModelUsage()) ? chatInfo.getSelectedModel(): agentConfig.getFixedModel();
        if (modelMark == null) {
            return null;
        }

        // 4. 获取历史对话
        List<Message> historyMessages = new ArrayList<>();
        if (agentConfig.getReferencedFunctions() != null && !agentConfig.getReferencedFunctions().isEmpty()) {
            historyMessages = messageStoreService.getHistoryMessages(chatInfo.getSessionId(), 100);
        }

        // 5. 背景提示词
        historyMessages.add(0, Message.builder().content(agentConfig.getBackgroundPrompt())
                .role(ChatMessageTypeEnum.SYSTEM.getCode()).build());

        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(100);

        historyMessages.stream().map(ChatServiceImpl::transMessage).forEach(chatMemory::add);

        // 6. 构建提示词
        String finalPrompt = buildFinalPrompt(chatInfo, agentConfig);

        // 7. 向量知识库检索
        List<PromptVectorConfig> vectorConfigList = agentConfigService.getVectorConfigById(chatInfo.getAgentId());
        if (vectorConfigList != null && !vectorConfigList.isEmpty()) {

            Map<String, String> ragResult = ragHelper.search(chatInfo.getUserMessage(), vectorConfigList);
            if (ragResult != null && !ragResult.isEmpty()) {

                for (Entry<String, String> entry : ragResult.entrySet()) {  
                    String key = entry.getKey();
                    String value = entry.getValue();
                    finalPrompt = finalPrompt.replace("$" + key + "$", value);
                }
            }
        }
        
        /*String knowledgeContent = "";
        if (agentConfig.getVectorStores() != null && !agentConfig.getVectorStores().isEmpty()) {
            try{
                knowledgeContent = vectorService.search(
                        agentConfig.getVectorStores(),
                        chatInfo.getUserMessage());
                finalPrompt = finalPrompt.replace("$ret_result$", knowledgeContent);
            }catch(Exception e){
                log.error("向量知识库检索失败", e);
            }
        }*/

        // 8. 处理函数调用
        String[] functionNames = {};
        if (chatInfo.getFunctionName() != null) {
            functionNames = new String[]{chatInfo.getFunctionName()};
        } else if (agentConfig.getReferencedFunctions() != null && !agentConfig.getReferencedFunctions().isEmpty()) {
            // 处理所有可用函数
            functionNames = agentConfig.getReferencedFunctions().split(",");
        }
        List<String> availableFunctions = new ArrayList<>();
        Map<ToolSpecification, ToolExecutor> toolsMap = new HashMap<>();
        for (String functionName : functionNames) {
            AiFunction function = aiFunctionProvider.getFunctionByName(functionName);
            if (function != null) {
                availableFunctions.add(functionName);
                toolsMap.put(Utils.convertToToolSpecification(function), aiFunctionProvider.toolExecutor);
            }
        }

        // 9. 调用模型服务
        return modelService.callSse(new ModelCallDTO(modelMark, agentConfig.getId(), finalPrompt, chatMemory,  toolsMap, agentConfig.getMaxtoken())); //TODO
    }

    @Override
    public Response<AiMessage> chatStream(AIChatInfo chatInfo, PrintWriter ioWriter) {
        String userId = chatInfo.getUserId();
        // 获取智能体信息
        AgentConfig agentConfig = agentConfigService.getById(chatInfo.getAgentId());
        if (agentConfig == null) {
            return null;
        }

        // 智能体鉴权
        if (!agentAuthService.checkAgentPermission(chatInfo.getAgentId(), userId)) {
            return null;
        }

        ChatDTO chat = null;
        String chatType = chatInfo.getType();
        if (StringUtils.isBlank(chatType)) {
            chat = makeChatCall(chatInfo, agentConfig);
        } else {
            switch (chatInfo.getType()) {
                case "chat":
                    chat = makeChatCall(chatInfo, agentConfig);
                break;
                case "writing":
                case "outline":
                    chat = makeWriteCall(chatInfo, agentConfig);
                    break;
            }
        }

        if (chat == null) {
            return null;
        }
        Long chatSessionId = chat.getChatSessionId();
        ModelCallDTO modelCallDTO = chat.getModelCallDTO();
        Response<AiMessage> chatResponse = null;
        try {
            // 调用模型服务
            chatResponse = modelService.callSse(modelCallDTO, ioWriter);

            if (chat.getOnComplete() != null) {
                chat.getOnComplete().accept(chatResponse);
            }
            String responseContent = chatResponse.content().text();

            aiChatLogService.logSuccess(AIChatLogDto.builder()
                .agentId(chatInfo.getAgentId())
                .userId(Util.getLongValue(userId))
                .modelName(modelCallDTO.getMark())
                .sessionId(chatInfo.getSessionId())
                .messageId(chatSessionId)
                .requestContent(chatInfo.getUserMessage())
                .responseContent(responseContent)
                .requestTime(chatInfo.getRequestTime())
                .responseTime(LocalDateTime.now())
                .build()
            );
        } catch (Exception e) {
            // 记录失败日志
            aiChatLogService.logFailure(AIChatLogDto.builder()
                .agentId(chatInfo.getAgentId())
                .userId(Util.getLongValue(userId))
                .modelName(modelCallDTO.getMark())
                .sessionId(chatInfo.getSessionId())
                .messageId(chatSessionId)
                .requestContent(chatInfo.getUserMessage())
                .requestTime(chatInfo.getRequestTime())
                .responseTime(LocalDateTime.now())
                .errorMessage(e.getMessage())
                .build()
            );
            throw new RuntimeException(e);
        }
        return chatResponse;
    }

    /**
     * 构建通用业务对话
     * @param chatInfo
     * @param agentConfig
     * @return
     */
    private ChatDTO makeChatCall(AIChatInfo chatInfo, AgentConfig agentConfig) {
        String userId = chatInfo.getUserId();
        // 确定使用的模型标识
        String modelMark = ModelConstants.MODEL_USAGE_DYNAMIC.equals(agentConfig.getModelUsage()) ? chatInfo.getSelectedModel(): agentConfig.getFixedModel();
        if (modelMark == null) {
            return null;
        }

        // 获取历史对话
        List<Message> historyMessages = new ArrayList<>();
        if (agentConfig.getReferencedFunctions() != null && !agentConfig.getReferencedFunctions().isEmpty()) {
            historyMessages = messageStoreService.getHistoryMessages(chatInfo.getSessionId(), 100);
        }

        // 背景提示词
        historyMessages.add(0, Message.builder().content(agentConfig.getBackgroundPrompt())
                                      .role(ChatMessageTypeEnum.SYSTEM.getCode()).build());

        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(100);

        historyMessages.stream().map(ChatServiceImpl::transMessage).forEach(chatMemory::add);

        String finalPrompt = buildFormattedPrompt(chatInfo, agentConfig);

        // 保存用户消息
        Message userMessage = Message.builder()
             .sessionId(chatInfo.getSessionId())
             .userId(userId)
             .content(finalPrompt)
             .role(ChatMessageTypeEnum.USER.getCode())
             .agentId(chatInfo.getAgentId())
             .createdAt(new Date()).build();

        Long chatSessionId = messageStoreService.saveMessage(userMessage);

        // 处理函数调用
        Map<ToolSpecification, ToolExecutor> toolsMap = buildTools(chatInfo, agentConfig);

        return ChatDTO.builder()
                   .chatSessionId(chatSessionId)
                   .modelCallDTO(new ModelCallDTO(modelMark, agentConfig.getId(), finalPrompt, chatMemory,  toolsMap, agentConfig.getMaxtoken()))
                   .onComplete((chatResponse) -> {
                       // 保存模型响应
                       String responseContent = chatResponse.content().text();
                       Message assistantMessage = Message.builder()
                         .sessionId(chatInfo.getSessionId())
                         .userId(userId)
                         .content(responseContent)
                         .role(ChatMessageTypeEnum.ASSISTANT.getCode())
                         .createdAt(new Date())
                         .build();

                       if(Objects.nonNull(chatResponse.content().toolExecutionRequests())) {
                           String callFuncNames = String.join(",", chatResponse.content().toolExecutionRequests().stream().map(ToolExecutionRequest::name).collect(Collectors.toSet()));
                           assistantMessage.setCallFuncName(callFuncNames);
                       }
                       messageStoreService.saveMessage(assistantMessage);
                   })
                   .build();
    }

    /**
     * 构建写作任务对话
     */
    private ChatDTO makeWriteCall(AIChatInfo chatInfo, AgentConfig agentConfig) {
        String userId = chatInfo.getUserId();
        // 确定使用的模型标识
        String modelMark = ModelConstants.MODEL_USAGE_DYNAMIC.equals(agentConfig.getModelUsage()) ? chatInfo.getSelectedModel(): agentConfig.getFixedModel();
        if (modelMark == null) {
            return null;
        }

        Message message = Message.builder().content(agentConfig.getBackgroundPrompt())
               .role(ChatMessageTypeEnum.SYSTEM.getCode()).build();
        ChatMemory chatMemory = MessageWindowChatMemory.withMaxMessages(100);
        chatMemory.add(transMessage(message));


        // 保存用户消息
        Message userMessage = Message.builder()
                                     .sessionId(chatInfo.getSessionId())
                                     .userId(userId)
                                     .content(chatInfo.getUserMessage())
                                     .role(ChatMessageTypeEnum.USER.getCode())
                                     .agentId(chatInfo.getAgentId())
                                     .createdAt(new Date())
                                     .build();

        Long chatSessionId = messageStoreService.saveMessage(userMessage);

        String finalPrompt = buildFormattedPrompt(chatInfo, agentConfig);

        // 处理函数调用
        Map<ToolSpecification, ToolExecutor> toolsMap = buildTools(chatInfo, agentConfig);

        return ChatDTO.builder()
                      .chatSessionId(chatSessionId)
                      .modelCallDTO(new ModelCallDTO(modelMark, agentConfig.getId(), finalPrompt, chatMemory,  toolsMap, agentConfig.getMaxtoken()))
                      .onComplete((chatResponse) -> {
                          // 保存模型响应
                          String responseContent = chatResponse.content().text();
                          Message assistantMessage = Message.builder()
                                                            .sessionId(chatInfo.getSessionId())
                                                            .userId(userId)
                                                            .content(responseContent)
                                                            .role(ChatMessageTypeEnum.ASSISTANT.getCode())
                                                            .createdAt(new Date())
                                                            .build();

                          if(Objects.nonNull(chatResponse.content().toolExecutionRequests())) {
                              String callFuncNames = String.join(",", chatResponse.content().toolExecutionRequests().stream().map(ToolExecutionRequest::name).collect(Collectors.toSet()));
                              assistantMessage.setCallFuncName(callFuncNames);
                          }
                          messageStoreService.saveWriteRecord(assistantMessage, chatInfo.getAdditionalParams());
                      })
                      .build();
    }

    /**
     * 提示词拼接
     * @param chatInfo
     * @param agentConfig
     * @return
     */
    private String buildFormattedPrompt(AIChatInfo chatInfo, AgentConfig agentConfig) {

        // 构建提示词
        String finalPrompt = buildFinalPrompt(chatInfo, agentConfig);

        // 向量知识库检索
        List<PromptVectorConfig> vectorConfigList = agentConfigService.getVectorConfigById(chatInfo.getAgentId());
        if (vectorConfigList != null && !vectorConfigList.isEmpty()) {

            Map<String, String> ragResult = ragHelper.search(chatInfo.getUserMessage(), vectorConfigList);
            if (ragResult != null && !ragResult.isEmpty()) {

                for (Entry<String, String> entry : ragResult.entrySet()) {  
                    String key = entry.getKey();
                    String value = entry.getValue();
                    finalPrompt = finalPrompt.replace("$" + key + "$", value);
                }
            }
        }
        return finalPrompt;
    }

    /**
     * 构建工具调用
     * @param chatInfo
     * @param agentConfig
     * @return
     */
    private Map<ToolSpecification, ToolExecutor> buildTools(AIChatInfo chatInfo, AgentConfig agentConfig) {

        String[] functionNames = {};
        if (chatInfo.getFunctionName() != null) {
            functionNames = new String[]{chatInfo.getFunctionName()};
        } else if (agentConfig.getReferencedFunctions() != null && !agentConfig.getReferencedFunctions().isEmpty()) {
            // 处理所有可用函数
            functionNames = agentConfig.getReferencedFunctions().split(",");
        }
        List<String> availableFunctions = new ArrayList<>();
        Map<ToolSpecification, ToolExecutor> toolsMap = new HashMap<>();
        for (String functionName : functionNames) {
            AiFunction function = aiFunctionProvider.getFunctionByName(functionName);
            if (function != null) {
                availableFunctions.add(functionName);
                toolsMap.put(Utils.convertToToolSpecification(function), aiFunctionProvider.toolExecutor);
            }
        }

        return toolsMap;
    }

    /**
     * 构建最终的提示词
     */
    private String buildFinalPrompt(AIChatInfo chatInfo, AgentConfig agentConfig) {
        StringBuilder promptBuilder = new StringBuilder();

        // 获取提示词模板
        String promptContent = Util.null2String(agentConfig.getPromptTemplate());

        // 处理占位符参数
        Map<String, Object> additionalParams = chatInfo.getAdditionalParams();

        if (additionalParams != null && !additionalParams.isEmpty()) {
            for (Entry<String, Object> entry : additionalParams.entrySet()) {
                promptContent = promptContent.replace("{" + Util.null2String(entry.getKey()) + "}",
                        Util.null2String(entry.getValue()));
            }
        }

        // 解析最终提示词内容
        promptBuilder.append(Util.isEmpty(promptContent) ? chatInfo.getUserMessage()
                : promptContent.replace("{user_message}", chatInfo.getUserMessage()));

        return promptBuilder.toString();
    }

    /*
     * 转换消息类型
     * 
     * @param message 消息
     * 
     * @return 转换后的消息
     */
    private static ChatMessage transMessage(Message message) {
        ChatMessage chatMessage = null;

        switch (ChatMessageTypeEnum.getByCode(message.getRole())) {
            case USER:
                chatMessage = UserMessage.from(message.getContent());
                break;
            case ASSISTANT:
                chatMessage = AiMessage.from(message.getContent());
                break;
            case SYSTEM:
                chatMessage = SystemMessage.from(message.getContent());
                break;
            case TOOL:
                chatMessage = ToolExecutionResultMessage.from(message.getCallFuncName(), message.getCallFuncName(),
                        message.getContent());
                break;
            case FUNCTION:
                chatMessage = ToolExecutionResultMessage.from(null, message.getCallFuncName(), message.getContent());
                break;
            default:
                log.warn("未知的消息类型: {}", message.getRole());
        }
        return chatMessage;
    }
}
